#include "data-collection-allocation.h"
#include "ns3/log.h"
#include "ns3/double.h"
#include "ns3/string.h"
#include "ns3/simulator.h"
#include "ns3/mobility-model.h"
#include "ns3/node.h"
#include <fstream>
#include <iomanip>

namespace ns3 {

NS_LOG_COMPONENT_DEFINE ("DataCollectionAllocation");

NS_OBJECT_ENSURE_REGISTERED (DataCollectionAllocation);

TypeId
DataCollectionAllocation::GetTypeId (void)
{
  static TypeId tid = TypeId ("ns3::DataCollectionAllocation")
    .SetParent<ChannelAllocation> ()
    .SetGroupName ("SAGINIoV")
    .AddConstructor<DataCollectionAllocation> ()
    .AddAttribute ("OutputDirectory",
                   "Directory where collected data will be stored",
                   StringValue ("./results/data-collection"),
                   MakeStringAccessor (&DataCollectionAllocation::SetOutputDirectory),
                   MakeStringChecker ())
  ;
  return tid;
}

DataCollectionAllocation::DataCollectionAllocation ()
  : m_outputDir ("./results/data-collection")
{
  NS_LOG_FUNCTION (this);
}

DataCollectionAllocation::~DataCollectionAllocation ()
{
  NS_LOG_FUNCTION (this);
  SaveResults ();
}

void
DataCollectionAllocation::SetOutputDirectory (std::string directory)
{
  NS_LOG_FUNCTION (this << directory);
  m_outputDir = directory;
  
  // 确保输出目录存在
  std::string mkdirCmd = "mkdir -p " + m_outputDir;
  system (mkdirCmd.c_str ());
}

std::map<uint32_t, int>
DataCollectionAllocation::AllocateChannels()
{
  NS_LOG_FUNCTION(this);
  
  std::map<uint32_t, int> allocation;

  // 获取当前时间
  Time now = Simulator::Now();
  double currentTime = now.GetSeconds();
  
  // 获取节点信息
  NodeContainer vehicles = m_nodeManager->GetVehicleNodes();
  NodeContainer groundStations = m_nodeManager->GetGroundStationNodes();
  NodeContainer satellites = m_nodeManager->GetSatelliteNodes();
  
  uint32_t numGroundStations = groundStations.GetN();
  uint32_t numSatellites = satellites.GetN();
  
  // 获取当前活动的车辆数据
  std::vector<VehicleData> vehicleData = m_nodeManager->GetVehicleDataAtTime(currentTime);
  
  NS_LOG_INFO("Collecting data for " << vehicleData.size() << " active vehicles at time " << currentTime << "s");

  // 为每个活跃的车辆收集数据
  for (const auto& vehicle : vehicleData) {
    // 找到对应的节点ID
    uint32_t vehicleNodeId = 0;
    
    // 遍历所有车辆节点，查找匹配vehicleId的节点
    for (uint32_t i = 0; i < vehicles.GetN(); i++) {
      std::string id = m_nodeManager->GetVehicleIdForNode(vehicles.Get(i)->GetId());
      if (id == vehicle.id) {
        vehicleNodeId = vehicles.Get(i)->GetId();
        break;
      }
    }
    
    // 如果找不到对应的节点，则继续下一个车辆
    if (vehicleNodeId == 0) {
      continue;
    }
    
    Ptr<Node> vehicleNode = nullptr;
    for (uint32_t i = 0; i < vehicles.GetN(); i++) {
      if (vehicles.Get(i)->GetId() == vehicleNodeId) {
        vehicleNode = vehicles.Get(i);
        break;
      }
    }
    
    if (!vehicleNode) {
      continue;
    }
    
    Ptr<MobilityModel> mobility = vehicleNode->GetObject<MobilityModel>();
    
    if (!mobility) {
      NS_LOG_WARN("Vehicle " << vehicleNodeId << " does not have mobility model");
      continue;
    }
    
    Vector position = mobility->GetPosition();
    Vector velocity = mobility->GetVelocity();
    double speed = sqrt(velocity.x * velocity.x + velocity.y * velocity.y + velocity.z * velocity.z);

    // 收集地面站信道数据
    if (m_allowGroundStationConnection) {
      for (uint32_t channelId = 0; channelId < numGroundStations; channelId++) {
        double sinr = CalculateSINR(vehicleNodeId, channelId);
        double throughput = CalculateThroughput(sinr, GROUND_STATION);
        double delay = CalculateDelay(vehicleNodeId, channelId);
        double interference = CalculateInterference(vehicleNodeId, channelId);
        double pathLoss = CalculatePathLoss(vehicleNodeId, channelId);
        
        // 创建数据记录
        DataRecord record;
        record.vehicleId = vehicleNodeId;
        record.channelId = channelId;
        record.throughput = throughput;
        record.delay = delay;
        record.interference = interference;
        record.pathLoss = pathLoss;
        record.sinr = sinr;
        record.time = currentTime;
        record.position = position;
        record.speed = speed;
        record.nodeType = "GroundStation";
        record.connectedNodeId = channelId + 667; // 假设地面站节点ID从667开始
        
        // 添加数据记录
        m_vehicleRecords[vehicleNodeId].push_back(record);
      }
    }
    
    // 收集卫星信道数据
    if (m_allowSatelliteConnection) {
      for (uint32_t j = 0; j < numSatellites; j++) {
        uint32_t channelId = numGroundStations + j;
        
        double sinr = CalculateSINR(vehicleNodeId, channelId);
        double throughput = CalculateThroughput(sinr, SATELLITE);
        double delay = CalculateDelay(vehicleNodeId, channelId);
        double interference = CalculateInterference(vehicleNodeId, channelId);
        double pathLoss = CalculatePathLoss(vehicleNodeId, channelId);
        
        // 创建数据记录
        DataRecord record;
        record.vehicleId = vehicleNodeId;
        record.channelId = channelId;
        record.throughput = throughput;
        record.delay = delay;
        record.interference = interference;
        record.pathLoss = pathLoss;
        record.sinr = sinr;
        record.time = currentTime;
        record.position = position;
        record.speed = speed;
        record.nodeType = "Satellite";
        record.connectedNodeId = channelId + 661; // 假设卫星节点ID从661开始
        
        // 添加数据记录
        m_vehicleRecords[vehicleNodeId].push_back(record);
      }
    }
  }

  // 每隔一段时间保存一次数据
  if (currentTime > 0 && ((uint32_t)currentTime % 60) == 0) {
    SaveResults();
  }
  
  // 返回空的分配结果，因为我们只需要收集数据，不需要实际分配
  return allocation;
}

void
DataCollectionAllocation::LogAllocationResults()
{
  NS_LOG_FUNCTION(this);
  
  NS_LOG_INFO("Data Collection algorithm doesn't perform actual allocation, only collects data");
  NS_LOG_INFO("Data is being saved to: " << m_outputDir);
}

void
DataCollectionAllocation::SaveResults()
{
  NS_LOG_FUNCTION(this);
  
  // 遍历每个车辆，保存其数据记录
  for (const auto& entry : m_vehicleRecords) {
    uint32_t vehicleId = entry.first;
    const std::vector<DataRecord>& records = entry.second;
    
    if (records.empty()) {
      continue;
    }
    
    // 创建CSV文件
    std::ostringstream filename;
    filename << m_outputDir << "/vehicle_" << vehicleId << ".csv";
    
    std::ofstream file(filename.str().c_str(), std::ios::out);
    if (!file) {
      NS_LOG_ERROR("Could not open file " << filename.str() << " for writing");
      continue;
    }
    
    // 写入CSV头部
    file << "vehicle_id,x,y,z,time,speed,channel_id,connected_node_id,node_type,sinr,throughput,delay,interference,path_loss" << std::endl;
    
    // 写入数据记录
    for (const DataRecord& record : records) {
      file << record.vehicleId << ","
           << record.position.x << ","
           << record.position.y << ","
           << record.position.z << ","
           << record.time << ","
           << record.speed << ","
           << record.channelId << ","
           << record.connectedNodeId << ","
           << record.nodeType << ","
           << record.sinr << ","
           << record.throughput << ","
           << record.delay << ","
           << record.interference << ","
           << record.pathLoss << std::endl;
    }
    
    file.close();
    NS_LOG_INFO("Saved data for vehicle " << vehicleId << " to " << filename.str());
  }
  
  // 清空记录
  m_vehicleRecords.clear();
}

} // namespace ns3
